import argparse

import numpy as np


def analyze_terminals(task_name):
    print(f"--- Analyzing task: {task_name} ---")
    try:
        npz_path = f"data_static/{task_name}.npz"
        with np.load(npz_path) as data:
            if "terminals" in data:
                terminals = data["terminals"]
                print("  Analysis based on raw 'terminals' from .npz:")

        # # Load dataset for the task
        # _, train_dataset, _ = ogbench.make_env_and_datasets(
        #     task_name,
        #     dataset_dir="data_static/",
        #     compact_dataset=True,
        # )

        # assert "valids" in train_dataset, "'valids' key not found in the dataset."

        # terminals = 1 - train_dataset["valids"]

        # if "valids" not in train_dataset:
        #     npz_path = f"data_static/{task_name}.npz"
        #     with np.load(npz_path) as data:
        #         if "terminals" in data:
        #             terminals = data["terminals"]
        #             print("  Analysis based on raw 'terminals' from .npz:")
        #         else:
        #             print(
        #                 f"Error: Neither 'valids' nor 'terminals' key found for task {task_name}"
        #             )
        #             return
        # else:
        #     terminals = 1 - train_dataset["valids"]
        #     print("  Analysis based on '1 - valids':")

        terminal_indices = np.where(terminals == 1)[0]

        if len(terminal_indices) < 2:
            print("  Not enough terminal points to calculate distances.")
            return

        distances = np.diff(terminal_indices)

        print(f"  Number of terminal sequences: {len(distances)}")
        print(f"  Min distance: {np.min(distances)}")
        print(f"  Max distance: {np.max(distances)}")
        print(f"  Mean distance: {np.mean(distances):.2f}")
        print(f"  Standard deviation: {np.std(distances):.2f}")

        if np.std(distances) > 1:
            print("  ==> Instability detected in action sequence length.")
        else:
            print("  ==> Action sequence length appears stable.")

    except Exception as e:
        print(f"Could not process task {task_name}. Error: {e}")


if __name__ == "__main__":
    # New list of tasks from the user's commands
    default_tasks = [
        "pointmaze-large-navigate-v0",
        "pointmaze-large-stitch-v0",
        "antmaze-large-navigate-v0",
        "antmaze-large-stitch-v0",
        "cube-double-play-v0",
        "puzzle-4x4-play-v0",
        "scene-play-v0",
    ]

    parser = argparse.ArgumentParser(
        description="Analyze terminal stability in OGBench datasets."
    )
    parser.add_argument(
        "--tasks",
        nargs="+",
        default=default_tasks,
        help="A list of ogbench task names to analyze.",
    )

    args = parser.parse_args()

    for task in args.tasks:
        analyze_terminals(task)
